import os
import shutil

import numpy as np

from constants import *
from option import parse_args


def first_five_frames_exist(real_video_crop):
    for i in range(5):
        image_name = '{}_0.png'.format(i * 10)
        if not os.path.exists(os.path.join(real_video_crop, image_name)):
            return False
    return True


def copy_five_frame(sampled_videos, original_path, manipulated_path, args):
    with open(os.path.join(args.dst_dir, 'real_fake_pairs.dat'), 'w') as f:
        for sampled_video in sampled_videos:
            sampled_video_split = sampled_video.split('_')
            real_video_crop = os.path.join(original_path, 'crops', sampled_video_split[0])
            fake_video_crop = os.path.join(manipulated_path, 'crops', sampled_video[:-4])
            # real, fake = sampled_video_split[0], sampled_video[:-4]
            if not first_five_frames_exist(real_video_crop):
                continue
            for i in range(5):
                image_name = '{}_0.png'.format(i * 10)
                if not os.path.exists(os.path.join(real_video_crop, image_name)):
                    continue

                dst_real_image = sampled_video_split[0] + '__' + image_name
                if os.path.exists(os.path.join(args.src_dir, dst_real_image)):
                    continue
                dst_fake_image = sampled_video[:-4] + '__' + image_name

                shutil.copy(os.path.join(real_video_crop, image_name), os.path.join(args.dst_dir, dst_real_image))
                shutil.copy(os.path.join(fake_video_crop, image_name), os.path.join(args.dst_dir, dst_fake_image))
                f.write('{},{}\n'.format(dst_real_image, dst_fake_image))


def get_test_videos(dataset_path, manipulated_path) -> list:
    test_videos = []
    lines = open(os.path.join(dataset_path, 'test.txt'), 'r').readlines()
    video_set = set([line.strip()[:-4] for line in lines])
    videos = os.listdir(os.path.join(manipulated_path, 'videos'))
    diffs_path = os.path.join(manipulated_path, 'diffs')
    for video in videos:
        video_split = video.split('_')
        assert len(video_split) == 2
        if video_split[0] not in video_set:
            continue
        # num_mask = len(os.listdir(os.path.join(diffs_path, video[:-4])))
        # if num_mask == 0:
        #     continue
        test_videos.append(video)
    return test_videos


def main():
    args = parse_args()
    args.src_dir = os.path.join(args.src_dir, 'FaceForensics++')
    if args.fake_type == DF:
        args.dst_dir = os.path.join(args.dst_dir, FACE_FORENSICS_DF)
    elif args.fake_type == F2F:
        args.dst_dir = os.path.join(args.dst_dir, FACE_FORENSICS_F2F)
    elif args.fake_type == FSH:
        args.dst_dir = os.path.join(args.dst_dir, FACE_FORENSICS_FSH)
    elif args.fake_type == FSW:
        args.dst_dir = os.path.join(args.dst_dir, FACE_FORENSICS_FSW)
    elif args.fake_type == NT:
        args.dst_dir = os.path.join(args.dst_dir, FACE_FORENSICS_NT)
    os.makedirs(args.dst_dir, exist_ok=True)

    original_path = os.path.join(args.src_dir, 'original_sequences', 'youtube', 'c23')
    manipulated_path = os.path.join(args.src_dir, 'manipulated_sequences', args.fake_type, 'c23')

    videos = get_test_videos(args.src_dir, manipulated_path)
    np.random.seed(111)
    # sampled_videos = np.random.choice(videos, 140)
    sampled_videos = videos
    sampled_videos.sort()
    # print(sampled_videos)

    copy_five_frame(sampled_videos, original_path, manipulated_path, args)


if __name__ == '__main__':
    # --src-dir /home/xinlin/data2 --dst-dir /home/xinlin/data2/test --fake-type Deepfakes
    # 'Deepfakes', 'Face2Face', 'FaceSwap', 'NeuralTextures', 'FaceShifter'
    main()
